#!/usr/bin/env python3
"""
Unitree RL Sim2Sim Demo Backend
提供MuJoCo仿真服务和API接口
"""

import asyncio
import json
import os
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional

import mujoco
import mujoco.viewer
import numpy as np
from fastapi import FastAPI, HTTPException, WebSocket
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, JSONResponse
from pydantic import BaseModel

# 配置路径
MODEL_PATH = Path("assets/unitree_go2/scene.xml")
REPORTS_DIR = Path("reports")
MODELS_DIR = Path("models")

# 创建必要的目录
REPORTS_DIR.mkdir(exist_ok=True)
MODELS_DIR.mkdir(exist_ok=True)

app = FastAPI(title="Unitree RL Sim2Sim Demo", version="1.0.0")

# 配置CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局变量
sim_data = {
    "is_running": False,
    "step_count": 0,
    "total_reward": 0.0,
    "velocity": 0.0,
    "energy": 0.0,
    "joint_positions": [],
    "joint_velocities": []
}

class SimulationRequest(BaseModel):
    model_path: Optional[str] = None
    policy_path: Optional[str] = None
    duration: float = 10.0
    render: bool = True

class ReportRequest(BaseModel):
    title: str
    description: str
    tags: List[str] = []

# MuJoCo仿真器
class MuJoCoSimulator:
    def __init__(self, model_path: str):
        self.model = mujoco.MjModel.from_xml_path(model_path)
        self.data = mujoco.MjData(self.model)
        self.viewer = None
        
    def reset(self):
        """重置仿真环境"""
        mujoco.mj_resetData(self.model, self.data)
        self.data.qpos[2] = 0.5  # 设置初始高度
        
    def step(self, action: np.ndarray):
        """执行一步仿真"""
        # 简单的控制策略示例
        self.data.ctrl[:] = action
        
        # 前进一步
        mujoco.mj_step(self.model, self.data)
        
        # 更新状态
        sim_data["step_count"] += 1
        sim_data["joint_positions"] = self.data.qpos.tolist()
        sim_data["joint_velocities"] = self.data.qvel.tolist()
        
        # 计算速度（简化的x方向速度）
        sim_data["velocity"] = np.linalg.norm(self.data.qvel[:3])
        
        # 计算能耗（简化的关节力矩平方和）
        sim_data["energy"] = np.sum(self.data.ctrl ** 2)
        
    def get_state(self) -> Dict:
        """获取当前状态"""
        return {
            "time": self.data.time,
            "qpos": self.data.qpos.tolist(),
            "qvel": self.data.qvel.tolist(),
            "ctrl": self.data.ctrl.tolist()
        }

# 初始化仿真器
simulator = None

@app.on_event("startup")
async def startup_event():
    """启动时初始化"""
    global simulator
    if MODEL_PATH.exists():
        simulator = MuJoCoSimulator(str(MODEL_PATH))
    else:
        print(f"警告: 模型文件 {MODEL_PATH} 不存在")

@app.get("/")
async def root():
    """根路径"""
    return {"message": "Unitree RL Sim2Sim Demo Backend", "status": "running"}

@app.get("/api/models")
async def list_models():
    """获取可用模型列表"""
    models = []
    if MODELS_DIR.exists():
        for file in MODELS_DIR.glob("*.pt"):
            models.append({
                "name": file.stem,
                "path": str(file),
                "size": file.stat().st_size
            })
    return {"models": models}

@app.post("/api/simulation/start")
async def start_simulation(request: SimulationRequest):
    """启动仿真"""
    global simulator
    
    if simulator is None:
        raise HTTPException(status_code=500, detail="仿真器未初始化")
    
    if sim_data["is_running"]:
        raise HTTPException(status_code=400, detail="仿真已在运行中")
    
    sim_data["is_running"] = True
    sim_data["step_count"] = 0
    sim_data["total_reward"] = 0.0
    
    simulator.reset()
    
    return {"status": "started", "message": "仿真已启动"}

@app.post("/api/simulation/stop")
async def stop_simulation():
    """停止仿真"""
    sim_data["is_running"] = False
    return {"status": "stopped", "message": "仿真已停止"}

@app.get("/api/simulation/status")
async def get_status():
    """获取仿真状态"""
    return {
        "is_running": sim_data["is_running"],
        "step_count": sim_data["step_count"],
        "velocity": sim_data["velocity"],
        "energy": sim_data["energy"],
        "joint_positions": sim_data["joint_positions"][:10],  # 只返回前10个关节
        "timestamp": datetime.now().isoformat()
    }

@app.post("/api/reports/generate")
async def generate_report(request: ReportRequest):
    """生成实验报告"""
    report = {
        "id": datetime.now().strftime("%Y%m%d_%H%M%S"),
        "title": request.title,
        "description": request.description,
        "tags": request.tags,
        "timestamp": datetime.now().isoformat(),
        "simulation_data": {
            "total_steps": sim_data["step_count"],
            "final_velocity": sim_data["velocity"],
            "total_energy": sim_data["energy"],
            "joint_positions": sim_data["joint_positions"],
            "joint_velocities": sim_data["joint_velocities"]
        }
    }
    
    # 保存报告
    report_path = REPORTS_DIR / f"report_{report['id']}.json"
    with open(report_path, 'w', encoding='utf-8') as f:
        json.dump(report, f, indent=2, ensure_ascii=False)
    
    return {"report_id": report["id"], "path": str(report_path)}

@app.get("/api/reports/{report_id}")
async def get_report(report_id: str):
    """获取实验报告"""
    report_path = REPORTS_DIR / f"report_{report_id}.json"
    if not report_path.exists():
        raise HTTPException(status_code=404, detail="报告不存在")
    
    return FileResponse(str(report_path))

@app.get("/api/reports")
async def list_reports():
    """获取报告列表"""
    reports = []
    for file in REPORTS_DIR.glob("report_*.json"):
        with open(file, 'r', encoding='utf-8') as f:
            report = json.load(f)
            reports.append({
                "id": report["id"],
                "title": report["title"],
                "timestamp": report["timestamp"]
            })
    
    return {"reports": reports}

@app.websocket("/ws/simulation")
async def websocket_endpoint(websocket: WebSocket):
    """WebSocket连接用于实时数据推送"""
    await websocket.accept()
    
    try:
        while True:
            if sim_data["is_running"]:
                # 发送实时数据
                await websocket.send_json({
                    "type": "status_update",
                    "data": {
                        "step_count": sim_data["step_count"],
                        "velocity": sim_data["velocity"],
                        "energy": sim_data["energy"],
                        "timestamp": datetime.now().isoformat()
                    }
                })
            await asyncio.sleep(0.1)  # 100ms更新一次
    except Exception as e:
        print(f"WebSocket连接错误: {e}")

async def run_simulation(duration: float = 10.0):
    """后台仿真任务"""
    if simulator is None:
        return
    
    start_time = asyncio.get_event_loop().time()
    
    while sim_data["is_running"]:
        if asyncio.get_event_loop().time() - start_time > duration:
            sim_data["is_running"] = False
            break
            
        # 简单的行走策略（正弦波控制）
        action = np.sin(sim_data["step_count"] * 0.1) * np.ones(simulator.model.nu)
        simulator.step(action)
        
        await asyncio.sleep(0.01)  # 100Hz仿真频率

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)
